import pandas as pd
import numpy as np
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.impute import SimpleImputer
from datetime import datetime, timedelta

# 1. 数据加载
def load_data():
    base_path = './new/'
    user_profile = pd.read_csv(f'{base_path}user_profile_table.csv')
    user_balance = pd.read_csv(f'{base_path}user_balance_table.csv')
    share_interest = pd.read_csv(f'{base_path}mfd_day_share_interest.csv')
    bank_shibor = pd.read_csv(f'{base_path}mfd_bank_shibor.csv')
    
    print("数据加载完成：")
    print(f"用户信息表大小: {user_profile.shape}")
    print(f"用户余额表大小: {user_balance.shape}")
    print(f"收益率表大小: {share_interest.shape}")
    print(f"拆借利率表大小: {bank_shibor.shape}")
    
    return user_profile, user_balance, share_interest, bank_shibor

# 2. 数据预处理
def preprocess_data(user_balance, share_interest, bank_shibor):
    user_balance['report_date'] = pd.to_datetime(user_balance['report_date'], format='%Y%m%d')
    share_interest['mfd_date'] = pd.to_datetime(share_interest['mfd_date'], format='%Y%m%d')
    bank_shibor['mfd_date'] = pd.to_datetime(bank_shibor['mfd_date'], format='%Y%m%d')
    
    daily_stats = user_balance.groupby('report_date').agg({
        'total_purchase_amt': 'sum',
        'total_redeem_amt': 'sum'
    }).reset_index()
    
    print("数据预处理完成：")
    print(f"每日统计数据大小: {daily_stats.shape}")
    
    return daily_stats

# 3. 特征工程
def create_features(daily_stats, share_interest, bank_shibor):
    df = daily_stats.merge(share_interest, left_on='report_date', right_on='mfd_date', how='left')
    df = df.merge(bank_shibor, left_on='report_date', right_on='mfd_date', how='left')
    
    df['date_numeric'] = df['report_date'].apply(lambda x: (x - pd.Timestamp('1970-01-01')).days)
    
    df['day_of_week'] = df['report_date'].dt.dayofweek
    df['day_of_month'] = df['report_date'].dt.day
    df['month'] = df['report_date'].dt.month
    
    for lag in [1, 2, 3, 7]:
        df[f'purchase_lag_{lag}'] = df['total_purchase_amt'].shift(lag)
        df[f'redeem_lag_{lag}'] = df['total_redeem_amt'].shift(lag)
    
    for window in [3, 7, 14]:
        df[f'purchase_ma_{window}'] = df['total_purchase_amt'].rolling(window=window).mean()
        df[f'redeem_ma_{window}'] = df['total_redeem_amt'].rolling(window=window).mean()
    
    print("特征工程完成：")
    print(f"特征数据大小: {df.shape}")
    print(f"特征列: {df.columns.tolist()}")
    
    return df

# 数据填充函数
def fill_missing_values(df, feature_columns):
    imputer = SimpleImputer(strategy='mean')
    df_filled = df.copy()
    df_filled[feature_columns] = imputer.fit_transform(df[feature_columns])
    return df_filled, imputer

# 4. 模型训练和预测
def train_and_predict(df):
    feature_columns = [col for col in df.columns if col not in 
                      ['report_date', 'mfd_date_x', 'mfd_date_y', 'total_purchase_amt', 'total_redeem_amt'] 
                      and df[col].dtype != 'datetime64[ns]']
    
    print("使用的特征列：", feature_columns)
    
    df_filled, imputer = fill_missing_values(df, feature_columns)
    
    purchase_model = GradientBoostingRegressor(
        n_estimators=100,
        learning_rate=0.1,
        max_depth=3,
        min_samples_split=2,
        min_samples_leaf=1,
        subsample=1.0,
        random_state=42
    )
    redeem_model = GradientBoostingRegressor(
        n_estimators=100,
        learning_rate=0.1,
        max_depth=3,
        min_samples_split=2,
        min_samples_leaf=1,
        subsample=1.0,
        random_state=42
    )
    
    df_clean = df_filled.dropna()
    print(f"清理后的数据大小: {df_clean.shape}")
    
    X = df_clean[feature_columns].astype(float)
    y_purchase = df_clean['total_purchase_amt']
    y_redeem = df_clean['total_redeem_amt']
    
    purchase_model.fit(X, y_purchase)
    redeem_model.fit(X, y_redeem)
    
    print("模型训练完成")
    
    return purchase_model, redeem_model, feature_columns, imputer

# 5. 生成预测结果
def generate_predictions(purchase_model, redeem_model, last_date, features, feature_columns, imputer):
    pred_dates = pd.date_range(start=last_date + timedelta(days=1), periods=30)
    
    predictions = []
    current_features = features[feature_columns].astype(float).values.reshape(1, -1)
    current_features = imputer.transform(current_features)
    
    # 获取历史数据的均值作为基准
    historical_purchase_mean = features['total_purchase_amt']
    historical_redeem_mean = features['total_redeem_amt']
    
    for date in pred_dates:
        purchase_pred = purchase_model.predict(current_features)[0]
        redeem_pred = redeem_model.predict(current_features)[0]
        
        # 更新特征值用于下一次预测
        for i, lag in enumerate([1, 2, 3, 7]):
            lag_purchase_idx = feature_columns.index(f'purchase_lag_{lag}')
            lag_redeem_idx = feature_columns.index(f'redeem_lag_{lag}')
            if i == 0:
                current_features[0, lag_purchase_idx] = purchase_pred
                current_features[0, lag_redeem_idx] = redeem_pred
        
        # 更新日期相关特征
        date_numeric_idx = feature_columns.index('date_numeric')
        day_of_week_idx = feature_columns.index('day_of_week')
        day_of_month_idx = feature_columns.index('day_of_month')
        month_idx = feature_columns.index('month')
        
        current_features[0, date_numeric_idx] = (date - pd.Timestamp('1970-01-01')).days
        current_features[0, day_of_week_idx] = date.dayofweek
        current_features[0, day_of_month_idx] = date.day
        current_features[0, month_idx] = date.month
        
        # 限制预测值在合理范围内并转换为整数
        purchase_pred = int(round(max(0, min(purchase_pred, historical_purchase_mean * 2))))
        redeem_pred = int(round(max(0, min(redeem_pred, historical_redeem_mean * 2))))
        
        # 使用严格的格式化字符串
        predictions.append(f"{date.strftime('%Y%m%d')},{purchase_pred},{redeem_pred}")
    
    # 直接写入文本文件
    output_path = './new/comp_predict_table.csv'
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(predictions))
    
    print("预测完成，前5行预测结果：")
    for line in predictions[:5]:
        print(line)
    
    return output_path

# 主函数
def main():
    try:
        print("开始加载数据...")
        user_profile, user_balance, share_interest, bank_shibor = load_data()
        
        print("\n开始数据预处理...")
        daily_stats = preprocess_data(user_balance, share_interest, bank_shibor)
        
        print("\n开始特征工程...")
        df = create_features(daily_stats, share_interest, bank_shibor)
        
        print("\n开始模型训练...")
        purchase_model, redeem_model, feature_columns, imputer = train_and_predict(df)
        
        print("\n开始生成预测...")
        last_date = df['report_date'].max()
        features = df.iloc[-1]
        output_path = generate_predictions(purchase_model, redeem_model, last_date, features, feature_columns, imputer)
        
        print(f"\n预测结果已保存到: {output_path}")
        
    except Exception as e:
        print(f"发生错误: {str(e)}")
        raise

if __name__ == "__main__":
    main()
